import torch
import pandas as pd
import torch.nn as nn
import torch.optim as optim
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score
import matplotlib.pyplot as plt
import logging as log

log.basicConfig(level=log.INFO)
log = log.getLogger()

# 1. 数据准备
# 读取CSV文件
data = pd.read_csv('diabetes.csv')
X = data.drop(columns=['Outcome'])  # 'Outcome' 是标签列的名称
y = data['Outcome']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 数据标准化
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

# 将数据转换为Tensor
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train.values, dtype=torch.long)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test.values, dtype=torch.long)

# 2. 模型定义
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.layer1 = nn.Linear(X_train.shape[1], 16)  # 特征数量
        self.layer2 = nn.Linear(16, 16)
        self.output = nn.Linear(16, 2)  # 2 classes

    def forward(self, x):
        x = torch.relu(self.layer1(x))
        x = torch.relu(self.layer2(x))
        x = self.output(x)
        return x

# 初始化模型
model = SimpleNN()

# 3. 训练准备
criterion = nn.CrossEntropyLoss()  # 交叉熵损失函数
optimizer = optim.Adam(model.parameters(), lr=0.01)  # Adam优化器
num_epochs = 20000

# 早停机制定义
class EarlyStopping:
    def __init__(self, patience=10, min_delta=0.001):
        self.patience = patience
        self.min_delta = min_delta
        self.best_loss = None
        self.counter = 0
        self.early_stop = False

    def __call__(self, val_loss):
        if self.best_loss is None:
            self.best_loss = val_loss
        elif val_loss > self.best_loss - self.min_delta:
            self.counter += 1
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_loss = val_loss
            self.counter = 0

# 创建 EarlyStopping 对象
early_stopping = EarlyStopping(patience=10)

# 保存损失值以绘制损失曲线
train_losses = []
val_losses = []

# 训练循环
for epoch in range(num_epochs):
    model.train()  # 切换到训练模式
    optimizer.zero_grad()  # 梯度清零

    # 前向传播
    outputs = model(X_train)
    loss = criterion(outputs, y_train)
    train_losses.append(loss.item())
    
    # 反向传播和优化
    loss.backward()
    optimizer.step()

    # 计算验证集损失
    model.eval()
    with torch.no_grad():
        val_outputs = model(X_test)
        val_loss = criterion(val_outputs, y_test)
        val_losses.append(val_loss.item())
    
    # 检查早停条件
    early_stopping(val_loss.item())
    if early_stopping.early_stop:
        log.info(f'Early stopping at epoch {epoch+1}')
        break

    # 打印损失
    if (epoch + 1) % 5 == 0:
        log.info(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}, Val Loss: {val_loss.item():.4f}')

# 绘制训练损失和验证损失曲线
plt.figure(figsize=(10, 6))
plt.plot(train_losses, label='Training Loss')
plt.plot(val_losses, label='Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Loss Curve with Early Stopping')
plt.legend()
plt.grid(True)
plt.show()

# 4. 模型评估
model.eval()  # 切换到评估模式
with torch.no_grad():
    # 预测训练集
    train_outputs = model(X_train)
    _, train_predicted = torch.max(train_outputs, 1)
    train_accuracy = accuracy_score(y_train, train_predicted)
    
    # 预测测试集
    test_outputs = model(X_test)
    _, test_predicted = torch.max(test_outputs, 1)
    test_accuracy = accuracy_score(y_test, test_predicted)

    log.info(f'Train Accuracy: {train_accuracy * 100:.2f}%')
    log.info(f'Test Accuracy: {test_accuracy * 100:.2f}%')

# 5. 预测新样本
new_sample = torch.tensor([[0.5, 1.2, 0.3, 0.8, 0.1, 0.4, 0.7, 0.6]], dtype=torch.float32)
with torch.no_grad():
    output = model(new_sample)
    _, predicted_class = torch.max(output, 1)
    log.info(f'Predicted Class for new sample: {predicted_class.item()}')
